from functools import cached_property, reduce
import re
from .seq import MTSequence, Sequence
from .gene import Region, Transcript, RegionType
from .snv import SNV, TranscriptAnnotation


class INS(SNV):
    @cached_property
    def s_start(self) -> int:
        return self.pos

    @cached_property
    def s_ref(self) -> Sequence | MTSequence:
        """standard ref"""
        return self.ref.__class__("")

    @cached_property
    def s_alt(self) -> Sequence | MTSequence:
        return self.alt[1:]

    def annotate_utr(self, transcript: Transcript, region1: Region, region2: Region, utr_size1: int, utr_size2: int) -> TranscriptAnnotation:
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region2 if transcript.strand == "+" else region1)
        distance1 = transcript.calc_utr_distance(region1, utr_size1, self.pos)
        distance2 = transcript.calc_utr_distance(region2, utr_size2, self.pos + 1)
        if transcript.strand == "+":
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region2)
            trans_anno.na_change = f"c.{distance1}_{distance2}ins{self.s_alt}"
        else:
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region1)
            trans_anno.na_change = f"c.{distance2}_{distance1}ins{self.s_alt.reverse_complement()}"
        return trans_anno

    def annotate_cds(self, transcript: Transcript, region1: Region, region2: Region, pos: int) -> TranscriptAnnotation:
        cdna = transcript.coding_dna
        new_cdna = cdna.substitute(pos, self.ref, self.alt)
        if transcript.strand == "+":
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region2)
        else:
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region1)
            cdna = cdna.reverse_complement()
            new_cdna = new_cdna.reverse_complement()
        start = cdna.difference_simple(new_cdna)
        alt = new_cdna[start - 1 : start + self.s_alt.size - 1]
        if start == 1:
            trans_anno.na_change = f"c.-1_1ins{alt}"
            trans_anno.region, trans_anno.region2 = ("UTR5", "UTR5") if transcript.has_utr5 else ("UpStream", "UpStream")
        elif start - 1 == cdna.size:
            trans_anno.na_change = f"c.{cdna.size}_{cdna.size}+1ins{alt}"
            trans_anno.region, trans_anno.region2 = ("UTR3", "UTR3") if transcript.has_utr3 else ("DownStream", "DownStream")
        else:
            unit = alt.dup_unit()
            if start > unit.size and unit == cdna[start - unit.size - 1 : start - 1]:
                trans_anno.na_change = f"c.{start-1}dup{unit}" if unit.size == 1 else f"c.{start-unit.size}_{start-1}dup{unit}"
                count = alt.size // unit.size
                if count > 1:
                    trans_anno.na_change += f"[{count}]"
            else:
                trans_anno.na_change = f"c.{start-1}_{start}ins{alt}"
            protein = cdna.protein()
            new_protein = new_cdna.protein()
            start, end1, end2 = protein.difference(new_protein)
            aa1 = protein[start - 1 : end1]
            aa2 = new_protein[start - 1 : end2]
            if alt.size % 3 == 0:
                trans_anno.event = "ins_inframe"
                if aa1.size == 0:
                    unit = aa2.dup_unit()
                    if start > unit.size and unit == protein[start - unit.size - 1 : start - 1]:
                        trans_anno.aa_change = f"p.{unit}{start-1}dup" if unit.size == 1 else f"p.{protein[start-unit.size-1]}{start-unit.size}_{protein[start-2]}{start-1}dup"
                        count = aa2.size // unit.size
                        if count > 1:
                            trans_anno.aa_change += f"[{count}]"
                    else:
                        if start == 1:
                            trans_anno.aa_change = f"p.{protein[start-1]}{start}-1_{protein[start-1]}{start}ins{aa2}"
                        elif start > protein.size:
                            trans_anno.aa_change = f"p.{protein[start-2]}{start-1}_{protein[start-2]}{start-1}+1ins{aa2}"
                        else:
                            trans_anno.aa_change = f"p.{protein[start-2]}{start-1}_{protein[start-1]}{start}ins{aa2}"
                elif aa1.size == 1:
                    trans_anno.aa_change = f"p.{aa1}{start-1}delins{aa2}"
                else:
                    trans_anno.aa_change = f"p.{aa1[0]}{start-1}_{aa1[-1]}{end1-1}delins{aa2}"
            else:
                if start < protein.size:
                    trans_anno.event = "ins_frameshift"
                    if aa2[0] == "*":
                        trans_anno.aa_change = f"p.{aa1[0]}{start}*"
                    else:
                        fsi = new_protein[start - 1 :].find("*")
                        fs = "?" if fsi == -1 else ("" if fsi == 0 else f"{fsi+1}")
                        trans_anno.aa_change = f"p.{aa1[0]}{start}{aa2[0]}fs*{fs}"
            if protein[0] == "M" and new_protein[0] != "M":
                trans_anno.event += "_startloss"
            if protein[-1] == "*" and new_protein.find("*") == -1:
                trans_anno.event += "_stoploss"
        return trans_anno

    def annotate_intron(self, transcript: Transcript, region1: Region, region2: Region, cds_size1: int, cds_size2: int) -> TranscriptAnnotation:
        distance1 = transcript.calc_cds_distance(region1, cds_size1, self.pos)
        distance2 = transcript.calc_cds_distance(region2, cds_size2, self.pos + 1)
        if transcript.strand == "+":
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region2)
            trans_anno.na_change = f"c.{distance1}_{distance2}ins{self.s_alt}"
        else:
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region1)
            trans_anno.na_change = f"c.{distance2}_{distance1}ins{self.s_alt.reverse_complement()}"
        values = list(map(int, reduce(lambda x, y: x + y, map(lambda z: re.findall(r"[\+\-](\d+)$", z), [distance1, distance2]))))
        if min(values) <= 2:
            trans_anno.event = "splicing"
            trans_anno.region = "splicing"
        return trans_anno

    def annotate_cross(self, transcript: Transcript, region1: Region, region2: Region, utr_size1: int, utr_size2: int, cds_size1: int, cds_size2: int):
        if (transcript.strand == "+" and region1.end < transcript.cds_start) or (transcript.strand == "-" and region1.end < transcript.cds_end):
            return self.annotate_cds(transcript, region1, region2, cds_size2)
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region2 if transcript.strand == "+" else region1)
        if transcript.strand == "+":
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region2)
            distance1 = transcript.calc_cds_distance(region1, cds_size1, self.pos)
            distance2 = transcript.calc_utr_distance(region2, utr_size2, self.pos + 1)
            trans_anno.na_change = f"c.{distance1}_{distance2}ins{self.s_alt}"
        else:
            trans_anno = TranscriptAnnotation.from_transcript(transcript, region1)
            distance1 = transcript.calc_utr_distance(region1, utr_size1, self.pos)
            distance2 = transcript.calc_cds_distance(region2, cds_size2, self.pos + 1)
            trans_anno.na_change = f"c.{distance2}_{distance1}ins{self.s_alt.reverse_complement()}"
        return trans_anno

    def annotate_mrna(self, transcript: Transcript) -> TranscriptAnnotation:
        region1, cds_size1, utr_size1 = transcript.find_region(self.pos)
        region2, cds_size2, utr_size2 = transcript.find_region(self.pos + 1)
        if region1 == region2:
            if region1.end < transcript.cds_start or region1.start > transcript.cds_end:
                # 调控区
                return self.annotate_utr(transcript, region1, region2, utr_size1, utr_size2)
            # 非调控区
            if region1.type == RegionType.CDS:
                # CDS
                return self.annotate_cds(transcript, region1, region2, cds_size1 + self.pos - region1.start + 1)
            # Intron
            return self.annotate_intron(transcript, region1, region2, cds_size1, cds_size2)
        # region2 != region2
        if region2.end < transcript.cds_start or region1.start >= transcript.cds_end:
            # 完全调控区
            return self.annotate_utr(transcript, region1, region2, utr_size1, utr_size2)
        if region1.start >= transcript.cds_start and region2.end <= transcript.cds_end:
            # 完全非调控区
            if (transcript.strand == "+" and region2.type == RegionType.CDS) or (transcript.strand == "-" and region1.type == RegionType.CDS):
                # CDS
                return self.annotate_cds(transcript, region1, region2, cds_size2)
            # Intron
            return self.annotate_intron(transcript, region1, region2, cds_size1, cds_size2)
        # 部分调控区
        return self.annotate_cross(transcript, region1, region2, utr_size1, utr_size2, cds_size1, cds_size2)

    def annotate_ncrna(self, transcript: Transcript) -> TranscriptAnnotation:
        region1, _, _ = transcript.find_region(self.pos)
        region2, _, _ = transcript.find_region(self.pos + 1)
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region2 if transcript.strand == "+" else region1)
        pos = self.pos - transcript.tx_start + 1
        dna = transcript.dna
        new_dna = dna.substitute(pos, self.ref, self.alt)
        if transcript.strand == "-":
            dna = dna.reverse_complement()
            new_dna = new_dna.reverse_complement()
        trans_anno.region = "ncRNA"
        start = dna.difference_simple(new_dna)
        alt = new_dna[start - 1 : start + self.s_alt.size - 1]
        if start == 1:
            trans_anno.na_change = f"n.-1_1ins{alt}"
        elif start - 1 == dna.size:
            trans_anno.na_change = f"n.{dna.size}_{dna.size}+1ins{alt}"
        else:
            unit = alt.dup_unit()
            trans_anno.na_change = f"n.{start-1}dup{unit}" if start > unit.size and unit == dna[start - unit.size - 1 : start - 1] else f"n.{start-1}_{start}ins{alt}"
        return trans_anno
